import os
import json
import yaml
from mindspore.train import Callback
from .dump.generate_dump_cls import Q
from .common import global_manage


global_manage.init()


class DumpCallback(Callback):
    def on_train_step_end(self, run_context):
        op_info = []
        cur_path = os.path.dirname(os.path.realpath(__file__))
        yaml_path = os.path.join(cur_path, "dump", "support_ops.yaml")
        with open(yaml_path, "r") as f:
            support_ops = yaml.safe_load(f).get("Primitive")
        for ops in support_ops:
            ops_cls = getattr(Q, ops)
            if not ops_cls.global_dict:
                continue
            op_info.append(ops_cls.global_dict)

        json_str = json.dumps(op_info, indent=4)
        save_path = os.path.join(global_manage.get_value("dump_path"), "ops_info.json")
        with open(save_path, "w") as f:
            f.write(json_str)


def dump_init(real_data=None, dump_path=None):
    if real_data:
        global_manage.set_value("real_data", real_data)
    if dump_path:
        global_manage.set_value("dump_path", dump_path)